{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "ab835cce",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['S数字:4564=18257825782578256E',\n",
       " 'S大写:5895=贰叁伍捌贰叁伍捌贰叁伍捌贰叁伍捌零E',\n",
       " 'S数字:3532=14129412941294128E',\n",
       " 'S字母:8984=etoeotoeotoeotoeyE',\n",
       " 'S小写:1412=五六四八五六四八五六四八五六四八EP',\n",
       " 'S数字:7764=31059105910591056E',\n",
       " 'S小写:8989=三五九五九五九五九五九五九五九五六E',\n",
       " 'S数字:6596=26386638663866384E',\n",
       " 'S字母:5759=wepeiepeiepeiepeyE',\n",
       " 'S数字:7736=30947094709470944E']"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%run common.ipynb\n",
    "\n",
    "[tokenizer.decode(i) for i in tokenizer.get_batch_data(prefix=True)[1]][:10]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "bcb3d7c8",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "model_ppo = ModelPPO(torch.load('gen.model'))\n",
    "model_ppo_ref = ModelPPO(torch.load('gen.model'))\n",
    "\n",
    "for i in model_ppo_ref.parameters():\n",
    "    i.requires_grad_(False)\n",
    "\n",
    "model_cls = torch.load('cls.model')\n",
    "model_cls.to(device)\n",
    "\n",
    "for i in model_cls.parameters():\n",
    "    i.requires_grad_(False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "c3e7a327",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([2, 0, 0, 3, 3, 1, 1, 3, 3, 2, 2, 3, 2, 0, 2, 3, 1, 0, 3, 0, 3, 0, 3, 3,\n",
       "         0, 2, 3, 0, 2, 2, 2, 0, 1, 0, 3, 1, 1, 2, 1, 0, 3, 3, 3, 3, 3, 2, 3, 3,\n",
       "         3, 2, 2, 2, 3, 2, 3, 2, 0, 2, 0, 2, 3, 1, 3, 3], device='cuda:0'),\n",
       " tensor([[ 1, 48, 47, 50,  6,  8,  7,  4, 51],\n",
       "         [ 1, 44, 45, 50,  6,  7,  4, 11, 51],\n",
       "         [ 1, 44, 45, 50, 11,  8, 11, 12, 51],\n",
       "         [ 1, 46, 47, 50, 10, 12, 11, 11, 51],\n",
       "         [ 1, 46, 47, 50,  8, 13, 13,  8, 51],\n",
       "         [ 1, 45, 49, 50, 13,  6, 11,  4, 51],\n",
       "         [ 1, 45, 49, 50,  5,  7, 11, 12, 51],\n",
       "         [ 1, 46, 47, 50,  9,  5, 12,  8, 51],\n",
       "         [ 1, 46, 47, 50,  9, 13,  7, 11, 51],\n",
       "         [ 1, 48, 47, 50,  6, 12, 12,  5, 51]], device='cuda:0'))"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "@torch.no_grad()\n",
    "def get_question():\n",
    "    label, question, _ = tokenizer.get_batch_data(prefix=True)\n",
    "    label = torch.LongTensor(label).to(device)\n",
    "\n",
    "    #只要问题部分,等号后面的内容切除\n",
    "    question = [i[:i.index(tokenizer.encoder['=']) + 1] for i in question]\n",
    "\n",
    "    #统一长度\n",
    "    lens = max([len(i) for i in question])\n",
    "    question = [[tokenizer.encoder['P']] * (lens - len(i)) + i\n",
    "                for i in question]\n",
    "\n",
    "    question = torch.LongTensor(question).to(device)\n",
    "\n",
    "    return label, question\n",
    "\n",
    "\n",
    "label, question = get_question()\n",
    "\n",
    "label, question[:10]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "b77e7d96",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[23, 21, 16, 14, 23, 21, 16, 14, 23, 21, 16, 14, 23, 21, 16, 14,  2,  0],\n",
       "        [42, 39, 34, 42, 42, 39, 34, 42, 42, 39, 34, 42, 42, 39, 34, 42,  2,  0],\n",
       "        [ 6, 13, 13,  7,  8, 13, 13,  5,  8, 13, 13,  5,  8, 13, 13,  5,  6,  2],\n",
       "        [ 6, 11,  9,  5,  4, 11,  9,  5,  4, 11,  9,  5,  4, 11,  9,  4, 12,  2],\n",
       "        [ 5, 13, 13, 11, 11, 13, 13, 11, 11, 13, 13, 11, 11, 13, 13, 11, 10,  2],\n",
       "        [37, 41, 34, 42, 37, 41, 34, 42, 37, 41, 34, 42, 37, 41, 34, 42, 34,  2],\n",
       "        [18, 17, 15, 16, 19, 17, 15, 16, 19, 17, 15, 16, 19, 17, 15, 16,  2,  0],\n",
       "        [36, 34, 41, 37, 42, 34, 41, 37, 42, 34, 41, 37, 42, 34, 41, 37, 40,  2],\n",
       "        [16, 17, 21, 19, 14, 17, 21, 19, 14, 17, 21, 19, 14, 17, 21, 18, 22,  2],\n",
       "        [35, 35, 39, 36, 39, 35, 39, 36, 39, 35, 39, 36, 39, 35, 39, 36, 38,  2]],\n",
       "       device='cuda:0')"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#如果question的长度确定,这里可以转换成批运算\n",
    "@torch.no_grad()\n",
    "def get_answer(question):\n",
    "    answer = generate(model_ppo.model_gen, question)\n",
    "\n",
    "    #裁剪,只要生成的部分\n",
    "    answer = answer[:, question.shape[1]:]\n",
    "\n",
    "    return answer\n",
    "\n",
    "\n",
    "answer = get_answer(question)\n",
    "\n",
    "answer[:10]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "720b4c0f",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([-1.3870, -1.4730,  3.6036, -1.4200, -1.7467, -0.8223,  3.5521,  3.4942,\n",
       "        -0.6207, -0.9800, -1.3754, -0.6482, -0.6081, -0.9560, -1.1991, -1.5163,\n",
       "        -1.4721, -0.5136,  3.2385, -0.7568, -1.9590, -1.7225, -0.8976, -0.7586,\n",
       "        -2.2114,  4.0351,  3.3967,  3.8929, -0.5727, -0.6243, -1.2016,  3.9614,\n",
       "        -0.4552, -0.9152,  3.4247, -0.8902, -0.5747, -1.6382, -0.8617,  3.7373,\n",
       "        -1.0044, -0.8499, -0.8065, -0.6340, -0.8241, -0.9206, -1.2685, -0.8730,\n",
       "        -0.7941, -1.1478, -0.5335, -0.7640, -1.9618, -0.4679,  3.3003, -0.8195,\n",
       "         3.8506, -1.7656,  4.1354, -0.9555,  3.6093,  3.5485, -0.4862, -0.7969],\n",
       "       device='cuda:0')"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "@torch.no_grad()\n",
    "def get_reward(question, answer, label):\n",
    "    input_ids = torch.cat((question, answer), 1)\n",
    "    attention_mask = (input_ids != tokenizer.encoder['P']).long()\n",
    "\n",
    "    with torch.no_grad():\n",
    "        logits = model_cls(input_ids=input_ids, attention_mask=attention_mask)\n",
    "\n",
    "    return logits.gather(1, label.reshape(-1, 1)).squeeze(1)\n",
    "\n",
    "\n",
    "reward = get_reward(question, answer, label)\n",
    "\n",
    "reward"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "a48cfa6f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'success'"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\"\"\"注释代码,可以不看\"\"\"\n",
    "\n",
    "\n",
    "#get_delta函数的原理解释,注释性代码\n",
    "#数学上和get_delta函数等价,但是运行效率低\n",
    "def get_delta_note(value, reward_kl):\n",
    "    #下一个词的value,减去当前词的value,相当于对value去基线,缩小数值方差\n",
    "    #每个词的value是相互独立的,前后词value的差,可以视为预测质量的衡量\n",
    "    value_next = torch.zeros_like(value)\n",
    "    value_next[:, :-1] = value[:, 1:].clone()\n",
    "\n",
    "    #在value中融合reward,kl\n",
    "    diff = reward_kl + value_next - value\n",
    "\n",
    "    #蒙特卡洛采样法估计Q函数,每个时刻的价值,等于后续所有价值的加权求和\n",
    "    #这里计算的其实就是adv\n",
    "    delta = []\n",
    "    for i in range(diff.shape[1]):\n",
    "        s = 0\n",
    "        for j in range(i, diff.shape[1]):\n",
    "            s += diff[:, j] * 0.95**(j - i)\n",
    "        delta.append(s)\n",
    "\n",
    "    return torch.stack(delta, dim=1)\n",
    "\n",
    "\n",
    "#只用一次循环就计算出delta,计算效率提高很多\n",
    "def get_delta_fast(value, reward_kl):\n",
    "    delta = []\n",
    "\n",
    "    for i in reversed(range(reward_kl.shape[1])):\n",
    "        value_next = 0\n",
    "        if i < reward_kl.shape[1] - 1:\n",
    "            value_next = value[:, i + 1]\n",
    "\n",
    "        diff = reward_kl[:, i] + value_next - value[:, i]\n",
    "\n",
    "        diff_last = 0\n",
    "        if len(delta):\n",
    "            diff_last = delta[-1]\n",
    "\n",
    "        delta.append(diff + 0.95 * diff_last)\n",
    "\n",
    "    return torch.stack(delta[::-1]).transpose(0, 1)\n",
    "\n",
    "\n",
    "#测试两个函数是等价的,误差是由于计算机精度导致的\n",
    "for _ in range(200):\n",
    "    value = torch.randn(64, 26)\n",
    "    reward_kl = torch.randn(64, 26)\n",
    "\n",
    "    assert (get_delta_note(value, reward_kl) -\n",
    "            get_delta_fast(value, reward_kl)).abs().max() < 1e-5\n",
    "'success'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "9d600b60",
   "metadata": {},
   "outputs": [],
   "source": [
    "from trl.core import clip_by_value, logprobs_from_logits, masked_mean, masked_whiten\n",
    "\n",
    "\n",
    "class PPOTrainer:\n",
    "\n",
    "    def __init__(self):\n",
    "        self.optimizer = torch.optim.Adam(model_ppo.parameters(), lr=1e-5)\n",
    "\n",
    "    def step(self, question, answer, reward):\n",
    "        with torch.no_grad():\n",
    "            #编码\n",
    "            token = [q.tolist() + a.tolist() for q, a in zip(question, answer)]\n",
    "            input_ids, attention_mask = tokenizer.batch_pad(token=token)\n",
    "            del token\n",
    "            input_ids = torch.LongTensor(input_ids).to(device)\n",
    "            attention_mask = torch.LongTensor(attention_mask).to(device)\n",
    "\n",
    "            #question和answer不需要内容,只需要长度信息即可\n",
    "            lens_q = [question.shape[1]] * len(question)\n",
    "            lens_a = []\n",
    "\n",
    "            for a in answer:\n",
    "                if tokenizer.encoder['E'] in a:\n",
    "                    lens_a.append(a.tolist().index(tokenizer.encoder['E']) + 1)\n",
    "                    continue\n",
    "                lens_a.append(len(a))\n",
    "\n",
    "            del question\n",
    "            del answer\n",
    "\n",
    "            #根据question计算answer的概率,并计算每个动作的分数\n",
    "            prob_log, value, mask = self.batched_forward_pass(\n",
    "                model_ppo, input_ids, attention_mask, lens_q, lens_a)\n",
    "\n",
    "            #使用ref模型计算概率,这是为了计算kl散度\n",
    "            prob_log_ref, _, _ = self.batched_forward_pass(\n",
    "                model_ppo_ref, input_ids, attention_mask, lens_q, lens_a)\n",
    "\n",
    "            #计算两份概率的kl散度,并融入reward\n",
    "            reward = self.compute_rewards(reward, prob_log, prob_log_ref, mask)\n",
    "\n",
    "            #计算delta和target,用于计算loss\n",
    "            value, delta, target = self.compute_advantages(value, reward, mask)\n",
    "\n",
    "        #每批数据循环N次模型\n",
    "        for _ in range(4):\n",
    "            #每次算一个数据\n",
    "            for i in range(len(input_ids)):\n",
    "                #重新计算概率和value\n",
    "                prob_log_new, value_new, _ = self.batched_forward_pass(\n",
    "                    model_ppo, input_ids[i].unsqueeze(0),\n",
    "                    attention_mask[i].unsqueeze(0), [lens_q[i]], [lens_a[i]])\n",
    "\n",
    "                #根据新旧概率求出变化率,进而求出loss\n",
    "                #根据target和value的差可以计算出另外一份loss\n",
    "                loss = self.get_loss(prob_log[i].unsqueeze(0),\n",
    "                                     value[i].unsqueeze(0), prob_log_new,\n",
    "                                     value_new, mask[i].unsqueeze(0),\n",
    "                                     delta[i].unsqueeze(0),\n",
    "                                     target[i].unsqueeze(0))\n",
    "\n",
    "                if not loss:\n",
    "                    continue\n",
    "\n",
    "                loss.backward()\n",
    "                #torch.nn.utils.clip_grad_norm_(model_ppo.parameters(), 1.0)\n",
    "                self.optimizer.step()\n",
    "                self.optimizer.zero_grad()\n",
    "\n",
    "    def batched_forward_pass(self, model, input_ids, attention_mask, lens_q,\n",
    "                             lens_a):\n",
    "        logits, value = model(input_ids=input_ids,\n",
    "                              attention_mask=attention_mask)\n",
    "\n",
    "        #取每个字的概率对数\n",
    "        prob_log = logprobs_from_logits(logits[:, :-1], input_ids[:, 1:])\n",
    "\n",
    "        #是预测结果并且不是PAD的位置是1\n",
    "        mask = torch.zeros_like(attention_mask)\n",
    "        mask[:, :-1] = attention_mask[:, 1:]\n",
    "        for i in range(len(input_ids)):\n",
    "            start = lens_q[i] - 1\n",
    "            end = start + lens_a[i]\n",
    "            mask[i, :start] = 0\n",
    "            mask[i, end:] = 0\n",
    "\n",
    "        #对最后一个字的预测没有意义,直接丢弃\n",
    "        value = value[:, :-1]\n",
    "        mask = mask[:, :-1]\n",
    "\n",
    "        return prob_log, value, mask\n",
    "\n",
    "    def compute_rewards(self, reward, prob_log, prob_log_ref, mask):\n",
    "        reward_kl = []\n",
    "\n",
    "        for i in range(len(reward)):\n",
    "            #求两份概率的kl散度\n",
    "            kl = self.get_kl(prob_log[i], prob_log_ref[i]) * -0.2\n",
    "\n",
    "            #把reward加在最后一个字的kl散度上\n",
    "            if (mask[i] == 0).all():\n",
    "                #print('all 0')\n",
    "                idx = 0\n",
    "            else:\n",
    "                idx = mask[i].nonzero()[-1].item()\n",
    "            kl[idx] += reward[i]\n",
    "\n",
    "            reward_kl.append(kl)\n",
    "\n",
    "        return torch.stack(reward_kl)\n",
    "\n",
    "    def compute_advantages(self, value, reward_kl, mask):\n",
    "        value = value * mask\n",
    "        reward_kl = reward_kl * mask\n",
    "\n",
    "        #这里计算delta的过程,可以看上面的注释.\n",
    "        delta = []\n",
    "        for i in reversed(range(reward_kl.shape[1])):\n",
    "            value_next = 0\n",
    "            if i < reward_kl.shape[1] - 1:\n",
    "                value_next = value[:, i + 1]\n",
    "\n",
    "            diff = reward_kl[:, i] + value_next - value[:, i]\n",
    "\n",
    "            diff_last = 0\n",
    "            if len(delta):\n",
    "                diff_last = delta[-1]\n",
    "\n",
    "            delta.append(diff + 0.95 * diff_last)\n",
    "\n",
    "        delta = torch.stack(delta[::-1]).transpose(0, 1)\n",
    "\n",
    "        #定义target,它估计了理想的value值\n",
    "        target = delta + value\n",
    "        delta = masked_whiten(delta, mask)\n",
    "\n",
    "        return value, delta, target\n",
    "\n",
    "    def get_loss(self, prob_log, value, prob_log_new, value_new, mask, delta,\n",
    "                 target):\n",
    "\n",
    "        #对数概率,相除变相减,取exp后还原为商,即两个模型输出logits的变化率\n",
    "        ratio = (prob_log_new - prob_log).exp()\n",
    "\n",
    "        #如果变化率太过于剧烈,可能是发生了震荡,跳过\n",
    "        if masked_mean(ratio, mask).item() > 10:\n",
    "            #print('skip', masked_mean(ratio, mask).item())\n",
    "            return None\n",
    "\n",
    "        #先算两个value的loss,简单的算mse loss就可以了\n",
    "        loss_vf1 = (value_new - target)**2\n",
    "        #数值裁剪,很显然是为了缓解自举\n",
    "        loss_vf2 = clip_by_value(value_new, value - 0.2, value + 0.2)\n",
    "        loss_vf2 = (loss_vf2 - target)**2\n",
    "        #两份loss取大的,还是为了缓解自举\n",
    "        loss_vf = 0.5 * masked_mean(torch.max(loss_vf1, loss_vf2), mask)\n",
    "\n",
    "        #计算ppo loss\n",
    "        loss_surr1 = -delta * ratio\n",
    "        #数值裁剪,很显然是为了缓解自举\n",
    "        loss_surr2 = -delta * ratio.clamp(0.8, 1.2)\n",
    "        loss_surr = masked_mean(torch.max(loss_surr1, loss_surr2), mask)\n",
    "\n",
    "        return loss_surr + 0.1 * loss_vf\n",
    "\n",
    "    def get_kl(self, a, b):\n",
    "        method = 'kl'\n",
    "\n",
    "        if method == 'kl':\n",
    "            return a - b\n",
    "\n",
    "        if method == 'abs':\n",
    "            return (a - b).abs()\n",
    "\n",
    "        if method == 'mse':\n",
    "            return (a - b).square() * 0.5\n",
    "\n",
    "        if method == 'full':\n",
    "            return torch.nn.functional.kl_div(a,\n",
    "                                              b,\n",
    "                                              log_target=True,\n",
    "                                              reduction='none')\n",
    "\n",
    "\n",
    "trainer = PPOTrainer()\n",
    "\n",
    "trainer.step(question, answer, reward)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "a1c70e8e",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 -0.1644705832004547\n",
      "S字母:8788= 三五一五五五一五五五一五五五一五二E -1.8778380155563354\n",
      "S大写:5525= 22102210221022100E -1.3465244770050049\n",
      "100 -0.13109271228313446\n",
      "S大写:5438= 21754175417541752E -1.740771770477295\n",
      "S字母:9446= 37787778777877784E -0.9766024351119995\n",
      "200 0.08248350769281387\n",
      "S字母:5045= 二〇一八二〇一八二〇一八二〇一八〇E -1.8033310174942017\n",
      "S字母:1891= 七五六四七五六四七五六四七五六四EP -1.385983943939209\n",
      "300 0.08983775973320007\n",
      "S大写:4288= 17153715371537152E -1.6275231838226318\n",
      "S字母:2708= qpieepitepiuepiewE 3.6285555362701416\n",
      "400 -0.011256426572799683\n",
      "S字母:3165= 一二六六一二六六一二六六一二六六〇E -1.7630295753479004\n",
      "S小写:6637= wyttpyttpyttpytriE -1.261555552482605\n",
      "500 1.0705180168151855\n",
      "S大写:7549= 30199019901990196E -1.7732865810394287\n",
      "S字母:7106= wirwyirwyirwyirwrE 3.213606595993042\n",
      "600 0.9927632808685303\n",
      "S小写:3522= qrpiorpiorpiorpiiE -1.0137674808502197\n",
      "S字母:9272= 37091709170917088E -0.6493400931358337\n",
      "700 1.253991723060608\n",
      "S字母:8986= 三五一四七五一四七五一四七五一四四E -1.7062468528747559\n",
      "S数字:2630= 10521052105210520E 3.4852147102355957\n",
      "800 1.9489874839782715\n",
      "S小写:6515= 二六〇六二六〇六二六〇六二六〇六〇E 3.5120785236358643\n",
      "S大写:7809= 三一二三九一二三九一二三九一二三六E -1.1651257276535034\n",
      "900 2.2832207679748535\n",
      "S小写:9063= 三六二五五六二五五六二五五六二五二E 3.869190216064453\n",
      "S小写:5056= 二〇二二六〇二二六〇二二六〇二二四E 3.8240041732788086\n",
      "1000 2.552072525024414\n",
      "S字母:6267= wtpuptpuptpuptpyiE 3.886591672897339\n",
      "S大写:7411= 二九六四六九六四六九六四六九六四四E -0.9950095415115356\n",
      "1100 1.9958045482635498\n",
      "S小写:2293= 九一七二三一七二三一七二三一七二EP 3.9680638313293457\n",
      "S大写:3712= 一四八四九四八四九四八四九四八四八E -1.0302479267120361\n",
      "1200 2.961498260498047\n",
      "S数字:1598= 6392639263926392EP 3.4238054752349854\n",
      "S小写:5128= 二〇五一四〇五一四〇五一四〇五一二E 3.8417978286743164\n",
      "1300 3.190735340118408\n",
      "S小写:1263= 五〇五二五〇五二五〇五二五〇五二EP 3.958763599395752\n",
      "S数字:9667= 38651865186518668E 3.9419965744018555\n",
      "1400 2.782565116882324\n",
      "S大写:9953= eoiqtoiqtoiqtoiqwE -0.39596277475357056\n",
      "S字母:4869= qoruuoruuoruuoruyE 3.7458271980285645\n",
      "1500 2.666975498199463\n",
      "S大写:7570= 三〇二八三〇二〇三〇二八三〇二八〇E -0.5756788849830627\n",
      "S数字:7847= 31311131113111308E 3.6758008003234863\n",
      "1600 2.0806758403778076\n",
      "S数字:5065= 20262026202620260E 3.6465954780578613\n",
      "S字母:7631= eptwuptwuptwuptwrE 3.6560962200164795\n",
      "1700 3.6355276107788086\n",
      "S大写:4187= 壹陆柒肆玖陆柒肆玖陆柒肆玖陆柒肆捌EPPPPPPP 3.55444073677063\n",
      "S小写:1774= 七〇九六七〇九六七〇九六七〇九六EPPPPPPPP 3.7898011207580566\n",
      "1800 3.7213969230651855\n",
      "S数字:7659= 30639063906390636E 3.8294143676757812\n",
      "S字母:6008= wrperrperrperrpe贰E 3.660592794418335\n",
      "1900 3.6552512645721436\n",
      "S字母:1361= trrrtrrrtrrrtrrrEP 3.942028760910034\n",
      "S小写:8038= 三二一五五二一五五二一五五二一五二E 4.008636951446533\n"
     ]
    }
   ],
   "source": [
    "for epoch in range(2000):\n",
    "    label, question = get_question()\n",
    "    answer = get_answer(question)\n",
    "    reward = get_reward(question, answer, label)\n",
    "\n",
    "    trainer.step(question, answer, reward)\n",
    "\n",
    "    if epoch % 100 == 0:\n",
    "        print(epoch, reward.mean().item())\n",
    "        for _, q, a, r in zip(range(2), question, answer, reward):\n",
    "            q = tokenizer.decode(q.tolist())\n",
    "            a = tokenizer.decode(a.tolist())\n",
    "            r = r.item()\n",
    "            print(q, a, r)\n",
    "\n",
    "model_ppo.to('cpu')\n",
    "torch.save(model_ppo, 'ppo.model')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:cuda117]",
   "language": "python",
   "name": "conda-env-cuda117-py"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
